#!/usr/bin/python

import sys
import matplotlib.pyplot as plt
import pylab
import numpy as np

filename = sys.argv[1]
f = open(filename, "rU")
fl = f.readlines()

tabfile = sys.argv[2]
sfh = open(tabfile, "rU")
sfl = sfh.readlines()

#outfile = sys.argv[3]

MIN = 1314756
MAX = 1338009

gene_starts = []
gene_stops = []
strands = []
p_arrows = []
m_arrows = []
p_ylocations = []
m_ylocations = []
ylocations = []

for line in sfl:
    l = line.split('\t')
    gene_starts.append(int(l[1]))
    gene_stops.append(int(l[2]))
    strands.append(int(l[3]))
    if int(l[3]) == 1:
        p_arrows.append(int(l[2]))
        p_ylocations.append(-3)
        ylocations.append(-3)
    else:
        m_arrows.append(int(l[1]))
        m_ylocations.append(-6)
        ylocations.append(-6)

print "Length of gene_starts: ", len(gene_starts)
print "Length of gene_stops: ", len(gene_stops)
print "Length of p_arrows: ", len(p_arrows)
print "Length of m_arrows: ", len(m_arrows)

#outfmt = sys.argv[3]

#action = sys.argv[4]

x1 = []
x2 = []
y1 = []
y2 = []

line = fl[4]
l = line.split('\t')
genome_size = int(l[7])

genome_coords = {}

for i in range(1, genome_size+1):
    genome_coords[i] = 0

for i, line in enumerate(fl[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    gs = int(l[7])
    identity = float(l[6])
    x1.append(start)
    x2.append(stop)
    y1.append(identity)
    y2.append(identity)
    for a in range(start, stop):
        genome_coords[a] = genome_coords[a] + 1

cov = []

for k, v in genome_coords.iteritems():
    cov.append(v)

gx = [1,genome_size]
gy = [10,10]

#a = np.array([x1, x2])
#b = np.array([y1, y2])

a = [x1, x2]
b = [y1, y2]

#xcoords = np.array(range(1, genome_size+1))
#ycoords = np.array(cov)

xcoords = range(1, genome_size+1)
#ycoords = cov

print "# of records in xcoords: ", len(xcoords)
print "# of records in cov: ", len(cov)

fig = plt.figure(1, figsize=(14,8))

j1 = [gene_starts, gene_stops]
k1 = [ylocations, ylocations]

plt.subplot(211)
plt.plot(xcoords, cov, color='green', linestyle='-')
plt.fill_between(xcoords, cov, facecolor='green', alpha=0.4)
plt.text(53330, 13000, "A", size=30)
plt.annotate('tmRNA', xy=(1000000, 5000), xycoords='data', xytext=(-30, 30), 
    textcoords='offset points', arrowprops=dict(arrowstyle="->"))
plt.annotate('rRNA operon', xy=(1708892, 5000), xycoords='data', xytext=(-50, 30), 
    textcoords='offset points', arrowprops=dict(arrowstyle="->"))
plt.annotate('rRNA operon', xy=(1963043, 10500), xycoords='data', xytext=(-30, 30), 
    textcoords='offset points', arrowprops=dict(arrowstyle="->"))
plt.annotate('rRNA operon', xy=(2851300, 10500), xycoords='data', xytext=(30, 0), 
    textcoords='offset points', arrowprops=dict(arrowstyle="->"))
plt.axvspan(MIN, MAX, facecolor='purple', alpha=0.3)
plt.axis([0, genome_size, 0, max(cov)+2])
#plt.title('Coverage of cDNA sequences for each position along the genome')
plt.xlabel('SGRA chromosome coordinates (bp)')
plt.ylabel('Counts')
plt.grid(True)

plt.subplot(212)
plt.plot(xcoords, cov, color='green', linestyle='-')
plt.plot(j1, k1, color='red', linestyle='-', lw=2.0)
plt.text(1315000, 200, "B", size=30)
plt.plot(p_arrows, p_ylocations, 'r>', mec='red')
plt.plot(m_arrows, m_ylocations, 'r<', mec='red')
plt.fill_between(xcoords, cov, facecolor='green', alpha=0.4)
plt.annotate('SGRA_1316', xy=(1332750, 100), xycoords='data', xytext=(-80, 30), 
    textcoords='offset points', arrowprops=dict(arrowstyle="->"))
plt.annotate('SGRA_1317', xy=(1334350, 85), xycoords='data', xytext=(-30, 50), 
    textcoords='offset points', arrowprops=dict(arrowstyle="->"))
#plt.axvspan(27544, 28873, facecolor='purple', alpha=0.3)
plt.axis([MIN, MAX, -10, max(cov[MIN:MAX])+2])
#plt.title('Coverage of cDNA sequences for each position along the genome')
#plt.title(r'$\mathit{S. grandis}\ $plasmid')
plt.xlabel('SGRA chromosome coordinates (bp)')
plt.ylabel('Counts')
plt.grid(True)

#if action == "show":
#    plt.show()
#elif action == "print":
#    plt.savefig(outfile, format=outfmt)

plt.show()
#plt.savefig(outfile, format=outfmt)

